import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import numpy as np
import pdb

class GlobalSTNNet(nn.Module):
    def __init__(self):
        super(GlobalSTNNet, self).__init__()
        # Spatial transformer localization-network
        self.localization = nn.Sequential(
            nn.Conv2d(3, 8, kernel_size=5, stride=2, padding=0),
            nn.MaxPool2d(2, stride=2),
            nn.PReLU(),
            nn.Conv2d(8, 16, kernel_size=3, stride=1, padding=0),
            nn.PReLU(),
            nn.Conv2d(16, 16, kernel_size=3, stride=1, padding=0),
            nn.MaxPool2d(2, stride=2),
            nn.PReLU(),
            nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=0),
            nn.PReLU(),
            nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=0),
            nn.MaxPool2d(2, stride=2),
            nn.PReLU(),
            nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=0),
            nn.PReLU(),
            #nn.Conv2d(48, 96, kernel_size=3, stride=1, padding=1),
            #nn.PReLU()
        )
        # Regressor for the 3 * 2 affine matrix
        self.fc_loc = nn.Sequential(
            nn.Linear(256, 32),
            nn.PReLU(),
            nn.Linear(32, 3 * 2)
        )

        self._initialize_weights()
        
        # Initialize the weights/bias with identity transformation
        #self.fc_loc[2].weight.data.zero_()
        #self.fc_loc[2].bias.data.copy_(torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float))

    def forward(self, x):
        # transform the input
        #return self.stn(x)
        #pdb.set_trace()
        xs = self.localization(x)
        xs = torch.flatten(xs, 1)
        out = self.fc_loc(xs)
        return  out 

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                n = m.weight.size(1)
                m.weight.data.normal_(0, 0.01)
                m.bias.data.zero_()

class LocalSTNNet(nn.Module):
    def __init__(self):
        super(LocalSTNNet, self).__init__()
        # Spatial transformer localization-network
        self.left_eye_net = self.make_backbone_layer()
        self.right_eye_net = self.make_backbone_layer()
        self.nose_net = self.make_backbone_layer()
        self.mouth_net = self.make_backbone_layer()
        # Regressor for the 3 * 2 affine matrix
        self.left_eye_loc = self.make_params_layer()
        self.right_eye_loc = self.make_params_layer()
        self.nose_loc = self.make_params_layer()
        self.mouth_loc = self.make_params_layer()

        self._initialize_weights()
        
        # Initialize the weights/bias with identity transformation
        #self.fc_loc[2].weight.data.zero_()
        #self.fc_loc[2].bias.data.copy_(torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float))

    def make_backbone_layer(self):
        backbone_net = nn.Sequential(
            nn.Conv2d(3, 8, kernel_size=5, stride=2, padding=0),
            nn.MaxPool2d(2, stride=2),
            nn.PReLU(),
            nn.Conv2d(8, 16, kernel_size=3, stride=1, padding=0),
            nn.PReLU(),
            nn.Conv2d(16, 16, kernel_size=3, stride=1, padding=0),
            nn.MaxPool2d(2, stride=2),
            nn.PReLU(),
            nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=0),
            nn.PReLU(),
            nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=0),
            nn.MaxPool2d(2, stride=2),
            nn.PReLU(),
            nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=0),
            nn.PReLU(),
            
        )
        return  backbone_net 

    def make_params_layer(self):
        params_net = nn.Sequential(
            nn.Linear(256, 32),
            nn.PReLU(),
            nn.Linear(32, 3 * 2) 
        )
        return  params_net 

    def forward(self, x):
        xs_1 = self.left_eye_net(x)
        xs_1 = torch.flatten(xs_1, 1)
        out_1 = self.left_eye_loc(xs_1)

        xs_2 = self.right_eye_net(x)
        xs_2 = torch.flatten(xs_2, 1)
        out_2 = self.right_eye_loc(xs_2)

        xs_3 = self.nose_net(x)
        xs_3 = torch.flatten(xs_3, 1)
        out_3 = self.nose_loc(xs_3)

        xs_4 = self.mouth_net(x)
        xs_4 = torch.flatten(xs_4, 1)
        out_4 = self.mouth_loc(xs_4)

        return  [out_1, out_2, out_3, out_4] 

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                n = m.weight.size(1)
                m.weight.data.normal_(0, 0.01)
                m.bias.data.zero_()


#################################################################
#################################################################
def conv_bn(inp, oup, stride):
    return nn.Sequential(
        nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
        nn.BatchNorm2d(oup),
        nn.ReLU(inplace=True)
    )


def conv_1x1_bn(inp, oup):
    return nn.Sequential(
        nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
        nn.BatchNorm2d(oup),
        nn.ReLU(inplace=True)
    )


class InvertedResidual(nn.Module):
    def __init__(self, inp, oup, stride, expand_ratio):
        super(InvertedResidual, self).__init__()
        self.stride = stride
        assert stride in [1, 2]

        hidden_dim = round(inp * expand_ratio)
        self.use_res_connect = self.stride == 1 and inp == oup

        if expand_ratio == 1:
            self.conv = nn.Sequential(
                # dw
                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
                nn.BatchNorm2d(hidden_dim),
                nn.ReLU(inplace=True),
                # pw-linear
                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
                nn.BatchNorm2d(oup),
            )
        else:
            self.conv = nn.Sequential(
                # pw
                nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False),
                nn.BatchNorm2d(hidden_dim),
                nn.ReLU(inplace=True),
                # dw
                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
                nn.BatchNorm2d(hidden_dim),
                nn.ReLU(inplace=True),
                # pw-linear
                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
                nn.BatchNorm2d(oup),
            )

    def forward(self, x):
        if self.use_res_connect:
            return x + self.conv(x)
        else:
            return self.conv(x)

class GlobalSTNMobileNetV2(nn.Module):
    def __init__(self, widen_factor=1.0): #, input_size=224
        super(GlobalSTNMobileNetV2, self).__init__()
        block = InvertedResidual
        input_channel = 8
        last_channel = 96
        interverted_residual_setting = [
            # t, c, n, s
            [1, 8, 1, 1],
            [6, 12, 2, 2],
            [6, 16, 2, 2],
            [6, 24, 3, 2],
            [6, 32, 3, 2],
            [6, 48, 3, 2],
            [6, 96, 1, 1],
        ]

        # building first layer
        # assert input_size % 32 == 0
        input_channel = int(input_channel * widen_factor)
        self.last_channel = int(last_channel * widen_factor) if widen_factor > 1.0 else last_channel
        self.features = [conv_bn(3, input_channel, 2)]
        # building inverted residual blocks
        for t, c, n, s in interverted_residual_setting:
            output_channel = int(c * widen_factor)
            for i in range(n):
                if i == 0:
                    self.features.append(block(input_channel, output_channel, s, expand_ratio=t))
                else:
                    self.features.append(block(input_channel, output_channel, 1, expand_ratio=t))
                input_channel = output_channel
        # building last several layers
        self.features.append(conv_1x1_bn(input_channel, self.last_channel))
        # make it nn.Sequential
        self.features = nn.Sequential(*self.features)
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))

        # building pts net
        self.fc_loc = nn.Sequential(
            nn.Linear(self.last_channel, 64),
            nn.PReLU(),
            nn.Linear(64, 64),
            nn.PReLU(),
            nn.Linear(64, 6)
        )


        self._initialize_weights()

    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        out = self.fc_loc(x)
        return out

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                n = m.weight.size(1)
                m.weight.data.normal_(0, 0.01)
                m.bias.data.zero_()

